package com.xzq.erp.ai.memory;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.xzq.erp.ai.dto.Msg;
import com.xzq.erp.constants.RedisConstants;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.Cursor;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ScanOptions;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import java.util.*;

@Component
public class RedisChatMemoryRepository implements ChatMemoryRepository {

    @Autowired
    private RedisTemplate<String, String>  redisTemplate;

    @Autowired
    private ObjectMapper objectMapper;

    private final String CHAT_KEY_PATTERN = "chat:memory:*";

    @Override
    public List<String> findConversationIds() {
        // 每个用户只能有10个会话
        return this.getAllChatMemoryKeys(CHAT_KEY_PATTERN, 5000);
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        String key = RedisConstants.buildChatMemoryKey(conversationId);
        List<String> list = redisTemplate.opsForList().range(key, 0, -1);
        if(Objects.isNull(list) || list.isEmpty()){
            return List.of();
        }
        // 倒序
        Collections.reverse(list);
        return list.stream().map(json -> {
            try {
                return objectMapper.readValue(json, Msg.class);
            } catch (JsonProcessingException e) {
                throw new RuntimeException(e);
            }
        }).map(Msg::toMessage).toList();
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        if(Objects.isNull(messages) ||  messages.isEmpty()){
            return;
        }
        List<String> list = messages.stream().map(Msg::new).map(msg -> {
                    try {
                        return objectMapper.writeValueAsString(msg);
                    } catch (JsonProcessingException e) {
                        throw new RuntimeException(e);
                    }
                }
        ).toList();
        String key = RedisConstants.buildChatMemoryKey(conversationId);
        redisTemplate.opsForList().leftPushAll(key,list);
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        redisTemplate.delete(RedisConstants.buildChatMemoryKey(conversationId));
    }

    public List<String> getAllChatMemoryKeys(String keyPattern, int count) {
        Set<String> keys = new HashSet<>();
        ScanOptions options = ScanOptions.scanOptions()
            .match(keyPattern)
            .count(count)  // 获取最多获取chatId的个数(与当前会话最大消息数有关)
            .build();
        Cursor<byte[]> cursor = redisTemplate.getConnectionFactory().getConnection().scan(options);
        while (cursor.hasNext()) {
            keys.add(new String(cursor.next()));
        }
        return new ArrayList<>(keys);
    }

}
